import fitz
from PIL import Image
import math
import os
import re


def pdf_to_long_image(pdf_path):
    doc = fitz.open(pdf_path)
    total_height = 0
    max_width = 0
    images = []
    for page_num in range(len(doc)):
        page = doc.load_page(page_num)
        pix = page.get_pixmap()
        img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
        images.append(img)
        total_height += img.height
        max_width = max(max_width, img.width)

    long_image = Image.new("RGB", (max_width, total_height))
    current_height = 0
    for img in images:
        long_image.paste(img, (0, current_height))
        current_height += img.height
    return long_image


def crop_horizontal_blank_space(image):
    width, height = image.size
    left = width
    right = 0
    for y in range(height):
        for x in range(width):
            pixel = image.getpixel((x, y))
            if pixel != (255, 255, 255):
                left = min(left, x)
                right = max(right, x)
    if left > right:
        return image
    return image.crop((left, 0, right + 1, height))


def is_solid_color(image):
    pixels = image.getdata()
    first_pixel = pixels[0]
    for pixel in pixels:
        if pixel != first_pixel:
            return False
    return True


def has_black_pixel(image):
    pixels = image.getdata()
    for pixel in pixels:
        if pixel == (0, 0, 0):
            return True
    return False


def get_next_index(output_folder):
    existing_indices = []
    for filename in os.listdir(output_folder):
        match = re.match(r'segment_(\d+)\.jpg', filename)
        if match:
            existing_indices.append(int(match.group(1)))
    if existing_indices:
        return max(existing_indices) + 1
    return 1


def crop_and_save(long_image, output_folder):
    width, height = long_image.size
    num_vertical_segments = math.ceil(height / 300)
    num_horizontal_segments = math.ceil(width / 400)
    segment_index = get_next_index(output_folder)
    for v in range(num_vertical_segments):
        start_y = v * 300
        end_y = min((v + 1) * 300, height)
        for h in range(num_horizontal_segments):
            start_x = h * 400
            end_x = min((h + 1) * 400, width)
            crop_img = long_image.crop((start_x, start_y, end_x, end_y))
            crop_width, _ = crop_img.size
            if (crop_width >= 20 or has_black_pixel(crop_img)) and not is_solid_color(crop_img):
                output_path = os.path.join(output_folder, f"segment_{segment_index}.jpg")
                crop_img.save(output_path, "JPEG")
                segment_index += 1


if __name__ == "__main__":
    pdf_path = "test.pdf"
    output_folder = "test1"
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    long_image = pdf_to_long_image(pdf_path)
    cropped_image = crop_horizontal_blank_space(long_image)
    crop_and_save(cropped_image, output_folder)
    print("处理完成！")
    
